package cn.caplike.demo.repository.javaagent.agentmain;

import javassist.*;

import java.io.IOException;
import java.lang.annotation.Annotation;
import java.lang.instrument.ClassFileTransformer;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.security.ProtectionDomain;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

public class TheClassFileTransformer implements ClassFileTransformer {

    public boolean supported(Class<?> clazz) {
        return Arrays.stream(clazz.getAnnotations()).anyMatch(annotation -> annotation.annotationType().getName().equals(Stub.class.getName()));
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classFileBuffer) {
        final ClassPool classPool = ClassPool.getDefault();
        classPool.insertClassPath(new LoaderClassPath(loader)); // Use host's classloader

        if (Objects.nonNull(classBeingRedefined) && supported(classBeingRedefined)) {
            System.out.println(" - Triggered by a redefine or retransform");

            try {
                final Annotation annotation = Arrays.stream(classBeingRedefined.getAnnotations()).filter(the -> the.annotationType().getName().equals(Stub.class.getName())).findAny().get();
                final Method method = annotation.getClass().getMethod("value");
                System.out.println(method.invoke(annotation));
            } catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
                throw new RuntimeException(e);
            }

            try {
                final CtClass ctClass = classPool.get(classBeingRedefined.getName());
                CtMethod setter = ctClass.getDeclaredMethod("setName", classPool.get(new String[]{"java.lang.String"}));
                setter.insertBefore("System.out.println(\"This method has been modified.\");");
                setter.insertAfter("super.put(\"name\", name);");
                try {
                    // addNewMethod(ctClass);

                    // todo: 将修改后的代码覆盖源文件
                    // ctClass.writeFile();
                    // ctClass.writeFile("BOOT-INF/classes/"/*protectionDomain.getCodeSource().getLocation().getPath()*/);
                    return ctClass.toBytecode();
                } catch (CannotCompileException e) {
                    System.err.println("Cannot compile due to `" + e.getReason() + "`");
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            } catch (NotFoundException | CannotCompileException ignored) {
                ignored.printStackTrace();
            }
        }

        return classFileBuffer;
    }

    private static void addNewMethod(CtClass ctClass) throws CannotCompileException {
        //新增一个方法
        CtMethod ctMethod = new CtMethod(CtClass.voidType, "newGenerated", new CtClass[]{}, ctClass);
        ctMethod.setModifiers(Modifier.PUBLIC);
        ctMethod.setBody("{System.out.println(\"This method is totally new.\");}");
        ctClass.addMethod(ctMethod);
    }
}
